package cn.sunxiang0918.tensorflowtest

import org.tensorflow.DataType
import org.tensorflow.Graph
import org.tensorflow.Output
import org.tensorflow.Tensor
import org.tensorflow.types.UInt8


/**
 * @author sun
 * @Date 2018/2/12 21:18
 * @version 1.0
 *
 */
class GraphBuilder(private val g: Graph) {

    fun div(x: Output<Float>, y: Output<Float>): Output<Float> {
        return binaryOp("Div", x, y)
    }

    fun <T> sub(x: Output<T>, y: Output<T>): Output<T> {
        return binaryOp("Sub", x, y)
    }

    fun <T> resizeBilinear(images: Output<T>, size: Output<Int>): Output<Float> {
        return binaryOp3("ResizeBilinear", images, size)
    }

    fun <T> expandDims(input: Output<T>, dim: Output<Int>): Output<T> {
        return binaryOp3("ExpandDims", input, dim)
    }

    fun <T, U> cast(value: Output<T>, type: Class<U>): Output<U> {
        val dtype = DataType.fromClass(type)
        return g.opBuilder("Cast", "Cast")
                .addInput(value)
                .setAttr("DstT", dtype)
                .build()
                .output(0)
    }

    fun decodeJpeg(contents: Output<String>, channels: Long): Output<UInt8> {
        return g.opBuilder("DecodeJpeg", "DecodeJpeg")
                .addInput(contents)
                .setAttr("channels", channels)
                .build()
                .output(0)
    }

    fun <T> constant(name: String, value: Any, type: Class<T>): Output<T> {
        Tensor.create(value, type).use { t ->
            return g.opBuilder("Const", name)
                    .setAttr("dtype", DataType.fromClass(type))
                    .setAttr("value", t)
                    .build()
                    .output(0)
        }
    }

    private fun constantInt(name: String, value: Any): Output<Int> {
        Tensor.create(value, java.lang.Integer::class.java).use { t ->
            return g.opBuilder("Const", name)
                    .setAttr("dtype", DataType.fromClass(java.lang.Integer::class.java))
                    .setAttr("value", t)
                    .build()
                    .output(0)
        }
    }

    private fun constantFloat(name: String, value: Any): Output<Float> {
        Tensor.create(value, java.lang.Float::class.java).use { t ->
            return g.opBuilder("Const", name)
                    .setAttr("dtype", DataType.fromClass(java.lang.Float::class.java))
                    .setAttr("value", t)
                    .build()
                    .output(0)
        }
    }

    fun constant(name: String, value: ByteArray): Output<String> {
        return this.constant(name, value, String::class.java)
    }

    fun constant(name: String, value: Int): Output<Int> {
        return this.constantInt(name, value)
    }

    fun constant(name: String, value: IntArray): Output<Int> {
        return this.constantInt(name, value)
    }

    fun constant(name: String, value: Float): Output<Float> {
        return this.constantFloat(name, value)
    }

    private fun <T> binaryOp(type: String, in1: Output<T>, in2: Output<T>): Output<T> {
        return g.opBuilder(type, type).addInput(in1).addInput(in2).build().output(0)
    }

    private fun <T, U, V> binaryOp3(type: String, in1: Output<U>, in2: Output<V>): Output<T> {
        return g.opBuilder(type, type).addInput(in1).addInput(in2).build().output(0)
    }
    
}